package com.serkol.ams.controller.api;

import com.serkol.ams.entity.LoginLog;
import com.serkol.constant.AuthConstant;
import com.serkol.openai.OpenAIService;
import com.serkol.openai.OpenDataAIService;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;


import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@RestController
public class ChatController {
    @Autowired
    private  OpenAiChatModel chatModel;

    @Autowired
    private OpenDataAIService openDataAIService;
    private final Map<Long, List<Message>> usersMsgs=new HashMap<>();
    private  List<SystemMessage> systemMessages=new ArrayList<>();

/*
    @Autowired
    public ChatController(OpenAiChatModel chatModel) {
        this.chatModel = chatModel;
    }
*/

    @GetMapping("/api/generate")
    public ChatResponse generate(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        LoginLog loginLog = AuthConstant.getLoginLog();
        if(systemMessages.isEmpty()) {
             systemMessages = openDataAIService.getSystemMessage(loginLog.getAgentId());
        }
        List<Message> messages = new ArrayList<>();
        if (loginLog == null) {
            // 用户未登录，返回空的聊天响应
            return new ChatResponse(List.of());
        } else {
            if (usersMsgs.containsKey(loginLog.getUserId())) {
                // 获取用户现有的消息列表
                List<Message> messages1 = usersMsgs.get(loginLog.getUserId());
                System.out.println("Len:" + messages1.size());
                messages.addAll(messages1);


                // 如果消息数量超过4条，删除最早的消息
                if (messages.size() > 4) {
                    messages = messages.subList(messages.size() - 4, messages.size());
                    messages.addAll(systemMessages);
                }
                messages.add(new UserMessage(message));
                usersMsgs.put(loginLog.getUserId(), messages);
            } else {
                // 如果用户没有聊天记录，创建新的消息列表
                //messages = List.of((Message) systemMessage,new UserMessage(message));
                messages.addAll(systemMessages);
                messages.add(new UserMessage(message));
                usersMsgs.put(loginLog.getUserId(), messages);
            }
        }

        // 创建Prompt并调用ChatModel
        Prompt prompt = new Prompt(messages);
        return this.chatModel.call(prompt);
    }


    @GetMapping("/api/generateStream")
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        Prompt prompt = new Prompt(new UserMessage(message));

        return this.chatModel.stream(prompt);
    }
    private final OpenAIService openAIService;

    public ChatController(OpenAIService openAIService) {
        this.openAIService = openAIService;
    }

    @GetMapping("api/chat")
    public Flux<String> chat(@RequestParam String message) {

        Flux<String> stringFlux = openAIService.streamChatCompletion(message);
        System.out.println("stringFlux Received message: " + message);
        stringFlux.subscribe(System.out::println);
        return stringFlux;
        //return Flux.just("Received message: " + message);
    }
}